# -*- coding:utf-8 -*-

# @Time    : 2019-01-09 13:52

# @Author  : Swing


import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
import os

os.environ['CUDA_VISIBLE_DEVICES'] = ''

print("Expected cross entropy loss if the model:")
print("- learns neither dependency:", -(0.625 * np.log(0.625) +
                                        0.375 * np.log(0.375)))
# Learns first dependency only ==> 0.51916669970720941
print("- learns first dependency:  ",
      -0.5 * (0.875 * np.log(0.875) + 0.125 * np.log(0.125))
      - 0.5 * (0.625 * np.log(0.625) + 0.375 * np.log(0.375)))
print("- learns both dependencies: ", -0.50 * (0.75 * np.log(0.75) + 0.25 * np.log(0.25))
      - 0.25 * (2 * 0.50 * np.log(0.50)) - 0.25 * (0))

num_steps = 10  # number of truncated backprop steps ('n' in the discussion above)
batch_size = 200
num_classes = 2
state_size = 16
learning_rate = 0.001


def gen_data(size=1000000):
    x = np.array(np.random.choice(2, size=(size,)))
    y = []
    for i in range(size):
        threshold = 0.5
        if x[i - 3] == 1:
            threshold += 0.5
        if x[i - 8] == 1:
            threshold -= 0.25
        if np.random.rand() > threshold:
            y.append(0)
        else:
            y.append(1)
    return x, np.array(y)


def gen_batch(raw_data, batch_size, num_steps):
    raw_x, raw_y = raw_data
    data_length = len(raw_x)

    batch_partition_length = data_length // batch_size
    data_x = np.zeros([batch_size, batch_partition_length], dtype=np.int32)
    data_y = np.zeros([batch_size, batch_partition_length], dtype=np.int32)

    for i in range(batch_size):
        data_x[i] = raw_x[batch_partition_length * i: batch_partition_length * (i + 1)]
        data_y[i] = raw_y[batch_partition_length * i: batch_partition_length * (i + 1)]

    epoch_size = batch_partition_length // num_steps

    for i in range(epoch_size):
        x = data_x[:, i * num_steps: (i + 1) * num_steps]
        y = data_y[:, i * num_steps: (i + 1) * num_steps]
        yield (x, y)


def gen_epochs(n, num_steps):
    for i in range(n):
        yield gen_batch(gen_data(), batch_size, num_steps)


tf.reset_default_graph()

xp = tf.placeholder(tf.int32, [batch_size, num_steps], name='input_placeholder')
yp = tf.placeholder(tf.int32, [batch_size, num_steps], name='labels_placeholder')
init_state = tf.zeros([batch_size, state_size])

xp_one_hot = tf.one_hot(xp, num_classes)
rnn_inputs = tf.unstack(xp_one_hot, axis=1)

cell = tf.nn.rnn_cell.BasicRNNCell(state_size)
rnn_outputs, final_state = tf.nn.static_rnn(cell, rnn_inputs, initial_state=init_state)
print(rnn_outputs)
print(final_state)

# with tf.variable_scope('rnn_cell'):
#     w = tf.get_variable('w', [num_classes + state_size, state_size])
#     b = tf.get_variable('b', [state_size], initializer=tf.constant_initializer(0.0))
#
#
# def rnn_cell(rnn_input, state):
#     with tf.variable_scope('rnn_cell', reuse=True):
#         w = tf.get_variable('w', [num_classes + state_size, state_size])
#         b = tf.get_variable('b', [state_size], initializer=tf.constant_initializer(0.0))
#     return tf.tanh(tf.matmul(tf.concat([rnn_input, state], 1), w) + b)
#
#
# state = init_state
# rnn_outputs = []
# for rnn_input in rnn_inputs:
#     state = rnn_cell(rnn_input, state)
#     rnn_outputs.append(state)
# final_state = rnn_outputs[-1]

with tf.variable_scope('softmax'):
    w = tf.get_variable('w', [state_size, num_classes])
    b = tf.get_variable('b', [num_classes], initializer=tf.constant_initializer(0.0))
logits = [tf.matmul(rnn_output, w) + b for rnn_output in rnn_outputs]
predictions = [tf.nn.softmax(logit) for logit in logits]

y_as_list = tf.unstack(yp, num=num_steps, axis=1)

losses = [tf.nn.sparse_softmax_cross_entropy_with_logits(labels=label, logits=logit) for logit, label in
          zip(logits, y_as_list)]
total_loss = tf.reduce_mean(losses)
train_step = tf.train.AdamOptimizer(learning_rate).minimize(total_loss)


def train_network(num_epochs, num_steps, state_size=4, verbose=True):
    with tf.Session() as sess:

        sess.run(tf.global_variables_initializer())
        training_losses = []

        for idx, epoch in enumerate(gen_epochs(num_epochs, num_steps)):
            training_loss = 0
            training_state = np.zeros((batch_size, state_size))

            if verbose:
                print("\nEPOCH", idx)

            for step, (x, y) in enumerate(epoch):
                tr_losses, training_loss_, training_state, _ = sess.run([losses, total_loss, final_state, train_step],
                                                                        feed_dict={
                                                                            xp: x,
                                                                            yp: y,
                                                                            init_state: training_state
                                                                        })

                training_loss += training_loss_

                if step % 100 == 0 and step > 0:
                    if verbose:
                        print("Average loss at step", step,
                              "for last 100 steps:", training_loss / 100)

                    training_losses.append(training_loss / 100)
                    training_loss = 0

    return training_losses


training_losses = train_network(10, num_steps, state_size=state_size)
plt.plot(training_losses)
plt.show()
